package com.kun.support.core;

import com.kun.constant.SecurityConstant;
import com.kun.service.KunUser;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 自定义jwt，将权限信息放至jwt中
 *
 * @author kun.li
 */
public class CustomeOauth2TokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {
    @Override
    public void customize(JwtEncodingContext context) {
        JwtClaimsSet.Builder claims = context.getClaims();
        if (OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType())) {
            // 检查登录用户信息是不是UserDetails，排除掉没有用户参与的流程
            if (context.getPrincipal().getPrincipal() instanceof UserDetails user) {
                // 获取申请的scopes
                Set<String> scopes = context.getAuthorizedScopes();
                // 获取用户的权限
                Collection<? extends GrantedAuthority> authorities = user.getAuthorities();
                // 提取权限并转为字符串
                Set<String> authoritySet = Optional.ofNullable(authorities).orElse(Collections.emptyList()).stream()
                        // 获取权限字符串
                        .map(GrantedAuthority::getAuthority)
                        // 去重
                        .collect(Collectors.toSet());
                // 合并scope与用户信息
                authoritySet.addAll(scopes);
                KunUser kunUser = (KunUser) user;
                // 将权限信息放入jwt的claims中（也可以生成一个以指定字符分割的字符串放入）
                setClaims(claims, authoritySet, kunUser, kunUser.getUserId(), kunUser.getTenantId(), kunUser.getHasAdmin());
            }
        }
        if (OAuth2TokenType.REFRESH_TOKEN.equals(context.getTokenType())){
            OAuth2Authorization authorization = context.getAuthorization();
            Map<String, Object> accessTokenClams = authorization.getAccessToken().getClaims();
            setClaims(claims, accessTokenClams.get(SecurityConstant.AUTHORITIES),
                    accessTokenClams.get(SecurityConstant.USERINFO),
                    accessTokenClams.get(SecurityConstant.USERID), accessTokenClams.get(SecurityConstant.TENANT_ID),
                    accessTokenClams.get(SecurityConstant.HAS_ADMIN));
        }

    }

    private void setClaims(JwtClaimsSet.Builder claims, Object authoritySet, Object userInfo, Object userId, Object tenantId, Object hasAdmin) {
        claims.claim(SecurityConstant.AUTHORITIES, authoritySet);
        claims.claim(SecurityConstant.USERINFO, userInfo);
        claims.claim(SecurityConstant.USERID, userId);
        claims.claim(SecurityConstant.TENANT_ID, tenantId);
        claims.claim(SecurityConstant.HAS_ADMIN, hasAdmin);
    }
}
